from typing import Annotated, Any
from typing_extensions import TypedDict
from langchain.pydantic_v1 import BaseModel, Field
import operator,re
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import create_react_agent
from ApiBase import apiBase
from ApiModels import VectInfo,apiModels


class PrimeInput(BaseModel):
    param: str = Field(default="1")

class LangGraphState(TypedDict):
    # The operator.add reducer fn makes this append-only
    aggregate: Annotated[list, operator.add]
    curNodeName: str    
    score: str
    output: str
    error: str

class LangGraphNode:
    def __init__(self, grp, vcts: list,agent_name: str):
        self.grp = grp
        # 第1个向量库存放
        self.vcts = vcts
        self.agent_name=agent_name
        self.agent=self.grp.react_create(self.vcts,self.agent_name)
    
    def __call__(self, state: LangGraphState) -> Any:
        # 取最后一个值
        quest=state['aggregate'][-1]
        # 能根据错误自动在向量库中 选择工具和选择prompt,动态创建agent
        fix_agent=None
        if quest['error']:
            fix_agent=self.grp.react_create(self.vcts,quest['error'])        
        else:
            fix_agent=self.agent
        #把返回格式 转成 输入格式        
        resp=self.grp.react_run(self.vcts,self.agent_name,fix_agent,quest['output'])
        resp["output"]={"messages": [("user", resp['output'])]}
        resp['curNodeName']=self.agent_name
        return {"aggregate": [resp]}
    
# 基于向量库调用工具
class ApiGrp():
    def __init__(self,apiTools):
        self.llm=apiTools.llm
        self.tol=apiTools.load_tol()
        self.vector=apiTools.load_vec()
    
    # 加载langgraph的node
    def langgraph_add_nodes(self, vcts,name_dict,edge_name,builder,node_tools):        
        if isinstance(edge_name, list):
            return
        if edge_name=="__start__":
            return
        if edge_name=="__end__":
            return
        if edge_name in name_dict:
            return
        name_dict[edge_name]=edge_name
        # prompt 让工具自己根据工具名称的关键字从向量库读取
        node_tools[edge_name]=LangGraphNode(self,vcts,edge_name)
        builder.add_node(edge_name, node_tools[edge_name])
    
    # 加载所有的hpl工具文件
    def load_json(self, vct_names:list,edges):
        node_tools={}
        name_dict={}
        vcts=apiModels.get_vects(vct_names)
        builder = StateGraph(LangGraphState)
        for edge in edges["edges"]:
            self.langgraph_add_nodes(vcts,name_dict,edge["from"],builder,node_tools)
            self.langgraph_add_nodes(vcts,name_dict,edge["to"],builder,node_tools)
            if edge["from"] == "__start__":
                edge["from"]=START
            if edge["to"] == "__end__":
                edge["to"]=END
            builder.add_edge(edge["from"], edge["to"])
        
        for edge in edges["conditional_edges"]:
            self.langgraph_add_nodes(vcts,name_dict,edge["from"],builder,node_tools)
            conds=self.tol.get_tool([edge["should_continue"]],type='tool-cond')
            # 要测试条件函数对象
            fun = conds[0]['fun']
            builder.add_conditional_edges(edge["from"], fun,edge["next"])
        graph = builder.compile()        
        return graph
    
    # 加载子节点
    def load_son(self,tab_name,parent,edges):
        parent_id=int(parent['id'])
        parent_name=parent['name']
        sql=f"select * from {tab_name} where parent_id={parent_id} and deleted <> 1"
        sub_nodes=apiBase.query_json(sql)
        if len(sub_nodes) == 0:
            # 最后一个节点默认加 end节点
            edge={}
            edge['from']=parent_name
            edge['to']=END
            edges['edges'].append(edge)
            return
        for sub in sub_nodes:
            # 把信息加载到向量库
            self.vector.train_QA(collect_name=sub['vct_name'],question=sub['name'],sql=sub['tools'],key=sub['name'],fun_name='tool-name')
            self.vector.train_QA(collect_name=sub['vct_name'],question=sub['name'],sql=sub['prompt'],key=sub['name'],fun_name='tool-prompt')
            if sub['cond']:
                self.vector.train_QA(collect_name=sub['vct_name'],question=sub['name'],sql=sub['cond'],key=sub['name'],fun_name='tool-cond')
                edges['conditional_edges'].append(sub['cond'])
            if sub['test']:
                self.vector.train_QA(collect_name=sub['vct_name'],question=sub['name'],sql=sub['test'],key=sub['name'],fun_name='tool-test')
            if sub['verify']:
                self.vector.train_QA(collect_name=sub['vct_name'],question=sub['name'],sql=sub['verify'],key=sub['name'],fun_name='tool-verify')
            edge={}
            # 如果是多个parent，就用;隔开
            if ";" in parent_name:
                edge['from']=parent_name.split(";")
            else:
                edge['from']=parent_name
            edge['to']=sub['name']
            edges['edges'].append(edge)
            self.load_son(tab_name,sub,edges)
        
    # record转成langgraph的json
    #agent表结构
    # grp_agent
    # 0、id
    # 1、agent的名称-agentA
    # 2、agent工具名称-database-sql-tools
    # 3、prompt内容-你是数据库专家
    # 4、test工具名称-''
    # 5、verify工具名称-''
    # 6、cond工具名称-'langgraph-sequential'
    # 5、parent_id-0
    def load_record(self, tab_name):
        edges={}
        edges['edges']=[]
        edges['conditional_edges']=[]
        parent={}
        parent['id']=0
        parent['name']='__start__'
        # 如果节点id为-1，表示这个
        self.load_son(tab_name,parent,edges)        
        return edges
    
    # 向量表根据工具名或prompt找到工具,创建react对象
    def react_create(self,vcts:list[VectInfo],agent_name):
        # 根据agent名称从向量库读取工具的名称
        tool=self.tol.get_tool(vcts,agent_name,"tool-name")
        # 根据tool_name根据关键字获取prompt，让工具自己根据工具名称的关键字从向量库读取
        prompts=self.vector.clts_query(vcts,agent_name,"tool-prompt",type=0)
        if len(prompts) == 0:
            return None
        # 只是取第1个
        usr_prompt=prompts[0]['answer']
        react = create_react_agent(self.llm, tools=tool, state_modifier=usr_prompt)            
        return react
    
    # 运行react对象
    def react_run(self,vcts:list[VectInfo],agent_name,agent,param):
        try:
            if not param:
                return "please input param"
            ret={}            
            msgs=agent.invoke(param)
            resp=msgs['messages'][-1].content
            pattern = r"```sql\n(.*?)```"
            final_code = re.findall(pattern, resp, re.DOTALL)[-1]
            ret["output"]=final_code
            ret['error']=self.tol.tol_chat(vcts,agent_name,ret['output'],'tool-test')
            if not ret['error']:
                ret['score']=self.tol.tol_chat(vcts,agent_name,ret['output'],'tool-verify')
            return ret
        except Exception as e:
           ret['error']=f'error :{str(e)}'
           return ret